import paddle
from einops import rearrange
from paddlets.models.forecasting.dl._crossformer.attention import AttentionLayer, TwoStageAttentionLayer


class DecoderLayer(paddle.nn.Layer):
    """
    The decoder layer of Crossformer, each layer will make a prediction at its scale
    """

    def __init__(self, seg_len, d_model, n_heads, d_ff=None, dropout=0.1,
        out_seg_num=10, factor=10):
        super(DecoderLayer, self).__init__()
        self.self_attention = TwoStageAttentionLayer(out_seg_num, factor,
            d_model, n_heads, d_ff, dropout)
        self.cross_attention = AttentionLayer(d_model, n_heads, dropout=dropout
            )
        self.norm1 = paddle.nn.LayerNorm(normalized_shape=d_model, epsilon=
            1e-05, weight_attr=None, bias_attr=None)
        self.norm2 = paddle.nn.LayerNorm(normalized_shape=d_model, epsilon=
            1e-05, weight_attr=None, bias_attr=None)
        self.dropout = paddle.nn.Dropout(p=dropout)
        self.MLP1 = paddle.nn.Sequential(paddle.nn.Linear(in_features=
            d_model, out_features=d_model), paddle.nn.GELU(), paddle.nn.
            Linear(in_features=d_model, out_features=d_model))
        self.linear_pred = paddle.nn.Linear(in_features=d_model,
            out_features=seg_len)

    def forward(self, x, cross):
        """
        x: the output of last decoder layer
        cross: the output of the corresponding encoder layer
        """
        batch = x.shape[0]
        x = self.self_attention(x)
        x = rearrange(x,
            'b ts_d out_seg_num d_model -> (b ts_d) out_seg_num d_model')
        cross = rearrange(cross,
            'b ts_d in_seg_num d_model -> (b ts_d) in_seg_num d_model')
        tmp = self.cross_attention(x, cross, cross)
        x = x + self.dropout(tmp)
        y = x = self.norm1(x)
        y = self.MLP1(y)
        dec_output = self.norm2(x + y)
        dec_output = rearrange(dec_output,
            '(b ts_d) seg_dec_num d_model -> b ts_d seg_dec_num d_model', b
            =batch)
        layer_predict = self.linear_pred(dec_output)
        layer_predict = rearrange(layer_predict,
            'b out_d seg_num seg_len -> b (out_d seg_num) seg_len')
        return dec_output, layer_predict


class Decoder(paddle.nn.Layer):
    """
    The decoder of Crossformer, making the final prediction by adding up predictions at each scale
    """

    def __init__(self, seg_len, d_layers, d_model, n_heads, d_ff, dropout,
        router=False, out_seg_num=10, factor=10):
        super(Decoder, self).__init__()
        self.router = router
        self.decode_layers = paddle.nn.LayerList()
        for i in range(d_layers):
            self.decode_layers.append(DecoderLayer(seg_len, d_model,
                n_heads, d_ff, dropout, out_seg_num, factor))

    def forward(self, x, cross):
        final_predict = None
        i = 0
        ts_d = x.shape[1]
        for layer in self.decode_layers:
            cross_enc = cross[i]
            x, layer_predict = layer(x, cross_enc)
            if final_predict is None:
                final_predict = layer_predict
            else:
                final_predict = final_predict + layer_predict
            i += 1
        final_predict = rearrange(final_predict,
            'b (out_d seg_num) seg_len -> b (seg_num seg_len) out_d', out_d
            =ts_d)
        return final_predict
